
from  Config.Config import *
import os
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import transforms

from Utils.ReaderProcess.ReadDict import ReadDict


class MyDataSet(Dataset):
    def __init__(self, data_dir,ClassesName, transform=None):
        self.ClassesName = ClassesName
        self.label_name = ReadDict.ReadModelClasses(self.ClassesName)

        self.data_info = self.get_img_info(data_dir)
        self.transform = transform

    def __getitem__(self, index):
        path_img, label = self.data_info[index]
        img = Image.open(path_img).convert('RGB')

        if self.transform is not None:
            img = self.transform(img)
        return img, label

    def __len__(self):
        return len(self.data_info)

    def get_img_info(self,data_dir):
        data_info = list()
        label_dict=ReadDict.ReadModelClasses(self.ClassesName)

        for root, dirs, _ in os.walk(data_dir): #
            # 遍历类别
            for sub_dir in dirs:
                img_names = os.listdir(os.path.join(root, sub_dir))
                img_names = list(filter(lambda x: x.endswith('.jpg'), img_names))
                # 遍历图片
                for i in range(len(img_names)):
                    img_name = img_names[i]
                    path_img = os.path.join(root, sub_dir, img_name)
                    label = label_dict[sub_dir]
                    data_info.append((path_img, int(label)))

        return data_info


if __name__ == '__main__':

    pass
    # train_dir=Data_Root+"\\"+Train
    # train_transform = transforms.Compose([
    #     transforms.Resize((32, 32)),
    #     transforms.ToTensor(),
    #
    # ])
    #
    # train_data = MyDataSet(data_dir=train_dir, transform=train_transform)
    # # 构建DataLoder
    # train_loader = DataLoader(dataset=train_data, batch_size=16)
    #
    # for data in train_data:
    #     img , label = data
    #     print(img.shape)



